import sys
import os
import requests
import json
import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import (Dense, LSTM, Dropout, BatchNormalization,
                                     Bidirectional, GRU, Conv1D, MaxPooling1D,
                                     Flatten, Layer, Input)
from tensorflow.keras.regularizers import l2
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.keras.utils import to_categorical
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QLabel, QPushButton,
                             QTextEdit, QSpinBox, QDoubleSpinBox, QProgressBar,
                             QMessageBox, QTabWidget, QGroupBox, QFormLayout, QComboBox)
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtGui import QIcon
import matplotlib

matplotlib.use('Qt5Agg')
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
from collections import Counter


def resource_path(relative_path):
    """ 解决打包后资源文件路径问题 """
    if hasattr(sys, '_MEIPASS'):
        return os.path.join(sys._MEIPASS, relative_path)
    return os.path.join(os.path.abspath("."), relative_path)


class AttentionLayer(Layer):
    def __init__(self, **kwargs):
        super(AttentionLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.W = self.add_weight(name='attention_weight',
                                 shape=(input_shape[-1], 1),
                                 initializer='random_normal',
                                 trainable=True)
        self.b = self.add_weight(name='attention_bias',
                                 shape=(input_shape[1], 1),
                                 initializer='zeros',
                                 trainable=True)
        super(AttentionLayer, self).build(input_shape)

    def call(self, x):
        e = tf.tanh(tf.matmul(x, self.W) + self.b)
        a = tf.nn.softmax(e, axis=1)
        output = x * a
        return tf.reduce_sum(output, axis=1)


def is_prime(n):
    """判断是否为质数"""
    if n <= 1:
        return False
    for i in range(2, int(n ** 0.5) + 1):
        if n % i == 0:
            return False
    return True


class DataFetcher(QThread):
    finished = pyqtSignal(str)
    progress = pyqtSignal(int)

    def run(self):
        try:
            url = 'http://www.cwl.gov.cn/cwl_admin/front/cwlkj/search/kjxx/findDrawNotice'
            params = {
                'name': 'ssq',
                'issueCount': '',
                'issueStart': '',
                'issueEnd': '',
                'dayStart': '',
                'dayEnd': '',
                'pageNo': '1',
                'pageSize': '9999',
                'week': '',
                'systemType': 'PC'
            }
            response = requests.get(url, params=params)
            jsondata = response.json()

            if jsondata['state'] == 0:
                data = []
                total = len(jsondata['result'])
                for i, item in enumerate(jsondata['result']):
                    date_str = item['date'].split('(')[0].strip()
                    blue_ball = item['blue']
                    red_balls = item['red'].split(',')
                    data.append([item['code'], date_str, int(red_balls[0]), int(red_balls[1]),
                                 int(red_balls[2]), int(red_balls[3]), int(red_balls[4]),
                                 int(red_balls[5]), int(blue_ball)])
                    self.progress.emit(int((i + 1) / total * 100))

                df = pd.DataFrame(data, columns=['期号', '日期', 'red1', 'red2', 'red3',
                                                 'red4', 'red5', 'red6', 'blue'])
                df.to_csv('data.csv', index=False, encoding='utf-8-sig')
                self.finished.emit(f"数据获取成功！共获取{len(data)}期数据。")
            else:
                self.finished.emit("数据获取失败：服务器返回错误状态。")
        except Exception as e:
            self.finished.emit(f"数据获取失败：{str(e)}")


class EnhancedPredictor(QThread):
    finished = pyqtSignal(str)
    progress = pyqtSignal(int)
    stats_ready = pyqtSignal(dict)
    model_trained = pyqtSignal(object)

    def __init__(self, train_ratio, epochs, batch_size, lookback, strategy):
        super().__init__()
        self.train_ratio = train_ratio
        self.epochs = epochs
        self.batch_size = batch_size
        self.lookback = lookback
        self.strategy = strategy.lower()
        self.red_stats = None
        self.blue_stats = None
        self.scaler = None
        self.model = None

    def get_model_type(self):
        """获取当前使用的模型类型"""
        if self.strategy == 'lstm':
            return "多层LSTM神经网络"
        elif self.strategy == 'gru':
            return "双向GRU神经网络"
        else:
            return "混合模型(CNN+LSTM+GRU+Attention)"

    def get_model_architecture(self):
        """获取模型架构描述"""
        if self.strategy == 'lstm':
            return "LSTM(512)→LSTM(256)→Dense(128)"
        elif self.strategy == 'gru':
            return "BiGRU(256)→GRU(128)→Dense(64)"
        else:
            return "Conv1D→BiLSTM→GRU→Attention→Dense"

    def calculate_enhanced_stats(self, data):
        """增强的统计分析"""
        stats = {}
        # 红球分析
        red_balls = []
        for i in range(1, 7):
            red_balls.extend(data[f'red{i}'].values)

        # 计算所有红球出现次数
        red_counts = pd.Series(red_balls).value_counts().sort_index()
        red_probs = (red_counts / red_counts.sum()).sort_values(ascending=False)

        # 近期分析（最近100期）
        recent_data = data.tail(100)
        recent_red = []
        for i in range(1, 7):
            recent_red.extend(recent_data[f'red{i}'].values)

        recent_red_counts = pd.Series(recent_red).value_counts()
        recent_red_probs = (recent_red_counts / recent_red_counts.sum()).sort_values(ascending=False)

        # 蓝球分析
        blue_counts = data['blue'].value_counts().sort_index()
        blue_probs = (blue_counts / blue_counts.sum()).sort_values(ascending=False)

        recent_blue_counts = recent_data['blue'].value_counts()
        recent_blue_probs = (recent_blue_counts / recent_blue_counts.sum()).sort_values(ascending=False)

        stats['red'] = {
            'all_time_top10': dict(list(red_probs.head(10).items())),
            'recent_top10': dict(list(recent_red_probs.head(10).items())),
            'all_time_sorted': dict(red_probs.sort_index())
        }

        stats['blue'] = {
            'all_time_top10': dict(list(blue_probs.head(10).items())),
            'recent_top10': dict(list(recent_blue_probs.head(10).items())),
            'all_time_sorted': dict(blue_probs.sort_index())
        }

        return stats

    def create_lstm_model(self, input_shape):
        """创建LSTM模型"""
        model = Sequential([
            Input(shape=input_shape),
            LSTM(512, return_sequences=True,
                 kernel_regularizer=l2(0.01), recurrent_regularizer=l2(0.01),
                 dropout=0.2, recurrent_dropout=0.2),
            BatchNormalization(),
            Bidirectional(LSTM(256, return_sequences=True)),
            BatchNormalization(),
            Dropout(0.3),
            LSTM(128),
            BatchNormalization(),
            Dropout(0.3),
            Dense(128, activation='relu', kernel_regularizer=l2(0.01)),
            Dense(64, activation='relu'),
            Dense(7, activation='sigmoid')
        ])
        model.compile(loss='mean_squared_error',
                      optimizer=Adam(learning_rate=0.001),
                      metrics=['mae'])
        return model

    def create_gru_model(self, input_shape):
        """创建GRU模型"""
        model = Sequential([
            Input(shape=input_shape),
            Bidirectional(GRU(256, return_sequences=True,
                              kernel_regularizer=l2(0.01),
                              recurrent_regularizer=l2(0.01),
                              dropout=0.2,
                              recurrent_dropout=0.2)),
            BatchNormalization(),
            GRU(128, return_sequences=True),
            BatchNormalization(),
            Dropout(0.3),
            GRU(64),
            BatchNormalization(),
            Dropout(0.3),
            Dense(64, activation='relu'),
            Dense(32, activation='relu'),
            Dense(7, activation='sigmoid')
        ])
        model.compile(loss='mean_squared_error',
                      optimizer=Adam(learning_rate=0.001),
                      metrics=['mae'])
        return model

    def create_hybrid_model(self, input_shape):
        """创建混合模型"""
        model = Sequential([
            Input(shape=input_shape),
            Conv1D(64, 3, activation='relu', padding='same'),
            MaxPooling1D(2),
            BatchNormalization(),
            Bidirectional(LSTM(256, return_sequences=True)),
            BatchNormalization(),
            GRU(128, return_sequences=True),
            AttentionLayer(),
            BatchNormalization(),
            Dropout(0.4),
            Dense(128, activation='relu'),
            Dense(64, activation='relu'),
            Dense(7, activation='sigmoid')
        ])
        model.compile(loss='mean_squared_error',
                      optimizer=Adam(learning_rate=0.001),
                      metrics=['mae'])
        return model

    def prepare_data(self, data):
        """数据预处理和特征工程"""
        data['date'] = pd.to_datetime(data['日期'], format='%Y-%m-%d', errors='coerce')
        if data['date'].isnull().any():
            data['date'] = pd.to_datetime(data['日期'].str.extract(r'(\d{4}-\d{2}-\d{2})')[0], format='%Y-%m-%d')

        # 添加时间特征
        data['year'] = data['date'].dt.year
        data['month'] = data['date'].dt.month
        data['day'] = data['date'].dt.day
        data['day_of_week'] = data['date'].dt.dayofweek
        data['day_of_year'] = data['date'].dt.dayofyear

        # 添加统计特征
        for i in range(1, 7):
            data[f'red{i}_rolling_mean_10'] = data[f'red{i}'].rolling(10).mean()
            data[f'red{i}_rolling_std_10'] = data[f'red{i}'].rolling(10).std()

        # 添加组合特征
        data['red_sum'] = data[[f'red{i}' for i in range(1, 7)]].sum(axis=1)
        data['red_odd_count'] = data[[f'red{i}' for i in range(1, 7)]].apply(lambda x: x % 2).sum(axis=1)
        data['red_prime_count'] = data[[f'red{i}' for i in range(1, 7)]].apply(lambda x: x.apply(is_prime)).sum(axis=1)

        # 添加滞后特征
        for lag in [1, 2, 3, 5, 10]:
            for i in range(1, 7):
                data[f'red{i}_lag{lag}'] = data[f'red{i}'].shift(lag)
            data[f'blue_lag{lag}'] = data['blue'].shift(lag)

        # 标准化
        cols_to_scale = [col for col in data.columns if col not in ['期号', '日期', 'date']]
        self.scaler = MinMaxScaler()
        scaled_data = self.scaler.fit_transform(data[cols_to_scale])

        return scaled_data, data[cols_to_scale].columns.tolist()

    def generate_sequences(self, data, lookback):
        """生成时间序列数据"""
        X, y = [], []
        for i in range(len(data) - lookback):
            X.append(data[i:i + lookback])
            y.append(data[i + lookback, :7])  # 只预测红球和蓝球
        return np.array(X), np.array(y)

    def enhanced_postprocessing(self, prediction):
        """改进的后处理方法"""
        red_pred = prediction[0][:6]
        blue_pred = prediction[0][6]

        # 处理红球
        red_balls = []
        for i in range(6):
            ball = int(round(red_pred[i] * 32 + 1))
            ball = max(1, min(33, ball))
            while ball in red_balls:
                ball = ball + 1 if ball < 33 else ball - 1
            red_balls.append(ball)

        red_balls = sorted(red_balls)

        # 处理蓝球
        blue_ball = int(round(blue_pred * 15 + 1))
        blue_ball = max(1, min(16, blue_ball))

        return red_balls, blue_ball

    def run(self):
        try:
            # 1. 数据加载
            self.progress.emit(5)
            data = pd.read_csv('data.csv')

            # 2. 统计分析
            self.progress.emit(10)
            stats = self.calculate_enhanced_stats(data)
            self.stats_ready.emit(stats)
            self.red_stats = stats['red']
            self.blue_stats = stats['blue']

            # 3. 数据预处理
            self.progress.emit(20)
            scaled_data, feature_names = self.prepare_data(data)

            # 4. 划分训练集和测试集
            self.progress.emit(30)
            train_size = int(len(scaled_data) * self.train_ratio)
            train_data = scaled_data[:train_size]
            test_data = scaled_data[train_size:]

            # 5. 生成序列数据
            self.progress.emit(40)
            X_train, y_train = self.generate_sequences(train_data, self.lookback)
            X_test, y_test = self.generate_sequences(test_data, self.lookback)

            # 6. 创建并训练模型
            self.progress.emit(50)
            if self.strategy == 'lstm':
                self.model = self.create_lstm_model((self.lookback, X_train.shape[2]))
            elif self.strategy == 'gru':
                self.model = self.create_gru_model((self.lookback, X_train.shape[2]))
            else:
                self.model = self.create_hybrid_model((self.lookback, X_train.shape[2]))

            # 自定义回调函数用于更新进度
            class ProgressCallback(tf.keras.callbacks.Callback):
                def __init__(self, progress_signal):
                    super().__init__()
                    self.progress_signal = progress_signal
                    self.epoch_count = 0

                def on_epoch_end(self, epoch, logs=None):
                    self.epoch_count += 1
                    progress = 50 + (self.epoch_count / self.params['epochs']) * 40
                    self.progress_signal.emit(int(progress))

            # 训练模型
            history = self.model.fit(
                X_train, y_train,
                epochs=self.epochs,
                batch_size=self.batch_size,
                validation_data=(X_test, y_test),
                callbacks=[
                    EarlyStopping(monitor='val_loss', patience=20),
                    ModelCheckpoint('best_model.h5', save_best_only=True),
                    ProgressCallback(self.progress)
                ],
                verbose=0
            )

            # 8. 预测下一期
            self.progress.emit(95)
            last_data = scaled_data[-self.lookback:]
            last_data = last_data[None, ...]
            prediction = self.model.predict(last_data)

            # 9. 后处理
            red_balls, blue_ball = self.enhanced_postprocessing(prediction)

            # 11. 生成最终结果
            result = self.generate_result(red_balls, blue_ball, stats)

            self.progress.emit(100)
            self.model_trained.emit(self.model)
            self.finished.emit(result)

        except Exception as e:
            self.finished.emit(f"预测失败：{str(e)}")

    def generate_result(self, red_balls, blue_ball, stats):
        """生成预测结果"""
        result = "=== 双色球优化预测结果 ===\n\n"
        result += f"预测号码：\n红球: {', '.join(map(str, red_balls))}\n蓝球: {blue_ball}\n\n"

        result += "=== 模型参数 ===\n"
        result += f"模型类型: {self.get_model_type()}\n"
        result += f"网络结构: {self.get_model_architecture()}\n"
        result += f"训练集比例: {int(self.train_ratio * 100)}%\n"
        result += f"训练轮次: {self.epochs}\n"
        result += f"批量大小: {self.batch_size}\n"
        result += f"回溯期数: {self.lookback}\n"

        return result


class StatsCanvas(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        self.fig.patch.set_facecolor('#F0F0F0')
        self.ax = self.fig.add_subplot(111)

        try:
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.rcParams['axes.unicode_minus'] = False
        except:
            try:
                plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
            except:
                pass

        super().__init__(self.fig)
        self.setParent(parent)

    def plot_stats(self, stats, ball_type):
        """绘制统计图表"""
        self.ax.clear()

        if ball_type == 'red':
            title = '红球出现频率TOP10'
            data = stats['red']['all_time_top10']
            color = '#E74C3C'
        else:
            title = '蓝球出现频率TOP10'
            data = stats['blue']['all_time_top10']
            color = '#3498DB'

        x = list(data.keys())
        y = [prob * 100 for prob in data.values()]

        bars = self.ax.bar(x, y, color=color, alpha=0.7)
        self.ax.set_title(title, fontsize=12, pad=15)
        self.ax.set_xlabel('号码', fontsize=10)
        self.ax.set_ylabel('出现概率(%)', fontsize=10)
        self.ax.grid(axis='y', linestyle='--', alpha=0.5)

        self.fig.tight_layout()
        self.draw()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("双色球预测系统 - 多维度神经网络模型 开发：Killerzeno")
        self.setGeometry(100, 100, 1200, 900)

        # 设置图标
        icon_path = resource_path('logo.ico')
        if os.path.exists(icon_path):
            self.setWindowIcon(QIcon(icon_path))

        self.stats_data = None
        self.model = None
        self.initUI()

    def initUI(self):
        main_widget = QWidget()
        layout = QVBoxLayout()

        # 创建标签页
        self.tabs = QTabWidget()

        # 第一页 - 预测功能
        predict_tab = QWidget()
        predict_layout = QVBoxLayout()

        # 数据获取部分
        data_group = QGroupBox("数据获取")
        data_layout = QVBoxLayout()

        self.fetch_btn = QPushButton("获取最新双色球数据")
        self.fetch_btn.clicked.connect(self.fetch_data)
        data_layout.addWidget(self.fetch_btn)

        self.fetch_progress = QProgressBar()
        data_layout.addWidget(self.fetch_progress)

        self.data_status = QTextEdit()
        self.data_status.setReadOnly(True)
        data_layout.addWidget(self.data_status)

        data_group.setLayout(data_layout)

        # 预测参数设置
        param_group = QGroupBox("模型参数设置")
        param_layout = QFormLayout()

        # 训练集比例
        self.train_ratio = QDoubleSpinBox()
        self.train_ratio.setRange(50, 95)
        self.train_ratio.setValue(75)
        self.train_ratio.setSingleStep(5)
        self.train_ratio.setSuffix("%")
        param_layout.addRow("训练集比例 (推荐70-80%):", self.train_ratio)

        # 训练轮次
        self.epochs = QSpinBox()
        self.epochs.setRange(50, 1000)
        self.epochs.setValue(200)
        param_layout.addRow("训练轮次 (推荐100-300):", self.epochs)

        # 批量大小
        self.batch_size = QSpinBox()
        self.batch_size.setRange(16, 128)
        self.batch_size.setValue(32)
        param_layout.addRow("批量大小 (推荐32-64):", self.batch_size)

        # 回溯期数
        self.lookback = QSpinBox()
        self.lookback.setRange(5, 30)
        self.lookback.setValue(15)
        param_layout.addRow("回溯期数 (推荐10-20):", self.lookback)

        # 预测策略
        self.strategy = QComboBox()
        self.strategy.addItems(['LSTM', 'GRU', '混合模型'])
        self.strategy.setCurrentText('LSTM')
        param_layout.addRow("预测策略:", self.strategy)

        param_group.setLayout(param_layout)

        # 预测按钮
        self.predict_btn = QPushButton("开始预测")
        self.predict_btn.clicked.connect(self.start_predict)

        # 预测进度
        self.predict_progress = QProgressBar()

        # 预测结果
        result_group = QGroupBox("预测结果")
        result_layout = QVBoxLayout()
        self.result_display = QTextEdit()
        self.result_display.setReadOnly(True)
        result_layout.addWidget(self.result_display)
        result_group.setLayout(result_layout)

        # 添加到预测页布局
        predict_layout.addWidget(data_group)
        predict_layout.addWidget(param_group)
        predict_layout.addWidget(self.predict_btn)
        predict_layout.addWidget(self.predict_progress)
        predict_layout.addWidget(result_group)
        predict_tab.setLayout(predict_layout)

        # 第二页 - 统计分析
        stats_tab = QWidget()
        stats_layout = QVBoxLayout()

        # 统计结果显示
        stats_text_group = QGroupBox("高频号码统计")
        stats_text_layout = QVBoxLayout()
        self.stats_display = QTextEdit()
        self.stats_display.setReadOnly(True)
        stats_text_layout.addWidget(self.stats_display)
        stats_text_group.setLayout(stats_text_layout)

        # 图表显示
        chart_group = QGroupBox("号码频率分布")
        chart_layout = QHBoxLayout()

        self.red_canvas = StatsCanvas(self, width=6, height=4)
        self.blue_canvas = StatsCanvas(self, width=6, height=4)

        chart_layout.addWidget(self.red_canvas)
        chart_layout.addWidget(self.blue_canvas)
        chart_group.setLayout(chart_layout)

        stats_layout.addWidget(stats_text_group)
        stats_layout.addWidget(chart_group)
        stats_tab.setLayout(stats_layout)

        # 添加标签页
        self.tabs.addTab(predict_tab, "预测")
        self.tabs.addTab(stats_tab, "统计分析")

        # 添加到主布局
        layout.addWidget(self.tabs)
        main_widget.setLayout(layout)
        self.setCentralWidget(main_widget)

        # 添加状态栏
        self.statusBar().showMessage("就绪")

    def fetch_data(self):
        if not self.check_internet_connection():
            QMessageBox.warning(self, "警告", "无法连接到互联网，请检查网络连接！")
            return

        self.fetch_btn.setEnabled(False)
        self.data_status.append("正在获取数据...")
        self.fetch_progress.setValue(0)

        self.fetcher = DataFetcher()
        self.fetcher.finished.connect(self.on_fetch_finished)
        self.fetcher.progress.connect(self.fetch_progress.setValue)
        self.fetcher.start()

    def check_internet_connection(self):
        """检查网络连接"""
        try:
            requests.get('http://www.baidu.com', timeout=5)
            return True
        except:
            return False

    def on_fetch_finished(self, message):
        self.fetch_btn.setEnabled(True)
        self.data_status.append(message)
        self.data_status.append("=" * 50)
        self.statusBar().showMessage("数据获取完成")

    def start_predict(self):
        if not os.path.exists('data.csv'):
            QMessageBox.warning(self, "警告", "请先获取数据！")
            return

        self.predict_btn.setEnabled(False)
        self.result_display.append("开始训练模型...")
        self.predict_progress.setValue(0)

        train_ratio = self.train_ratio.value() / 100
        epochs = self.epochs.value()
        batch_size = self.batch_size.value()
        lookback = self.lookback.value()
        strategy = self.strategy.currentText().lower()

        self.predictor = EnhancedPredictor(train_ratio, epochs, batch_size, lookback, strategy)
        self.predictor.finished.connect(self.on_predict_finished)
        self.predictor.progress.connect(self.predict_progress.setValue)
        self.predictor.stats_ready.connect(self.update_stats)
        self.predictor.model_trained.connect(self.set_model)
        self.predictor.start()

    def on_predict_finished(self, result):
        self.predict_btn.setEnabled(True)
        self.result_display.setPlainText(result)
        self.statusBar().showMessage("预测完成")

    def set_model(self, model):
        self.model = model

    def update_stats(self, stats):
        self.stats_data = stats

        # 更新统计文本显示
        text = "红球历史高频TOP10:\n"
        for i, (num, prob) in enumerate(stats['red']['all_time_top10'].items()):
            text += f"{i + 1}. 号码 {num}: {prob * 100:.2f}%\n"

        text += "\n红球近期高频TOP10:\n"
        for i, (num, prob) in enumerate(stats['red']['recent_top10'].items()):
            text += f"{i + 1}. 号码 {num}: {prob * 100:.2f}%\n"

        text += "\n蓝球历史高频TOP10:\n"
        for i, (num, prob) in enumerate(stats['blue']['all_time_top10'].items()):
            text += f"{i + 1}. 号码 {num}: {prob * 100:.2f}%\n"

        text += "\n蓝球近期高频TOP10:\n"
        for i, (num, prob) in enumerate(stats['blue']['recent_top10'].items()):
            text += f"{i + 1}. 号码 {num}: {prob * 100:.2f}%\n"

        self.stats_display.setPlainText(text)

        # 更新图表
        self.red_canvas.plot_stats(stats, 'red')
        self.blue_canvas.plot_stats(stats, 'blue')


if __name__ == "__main__":
    app = QApplication(sys.argv)
    icon_path = resource_path('logo.ico')
    if os.path.exists(icon_path):
        app.setWindowIcon(QIcon(icon_path))
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())